module Parser(parse, parseTest) where

import Prelude hiding (div, not)
import qualified Prelude
import Language.Python.Version2.Parser
import qualified Language.Python.Common as P
import Language.Python.Common.Pretty(prettyText)
import Types
import Misc(removeGroups, varsStmt)


parseTest :: String -> IO ()
parseTest filename = do
    contents <- readFile filename
    case parseModule contents filename of
        Left err -> error (prettyText err)
        Right (m, _) -> print m


parse :: String -> IO Program
parse filename = do
    contents <- readFile filename
    case parseModule contents filename of
        Left err -> error (prettyText err)
        Right (m, _) -> return (convModule m)


createMainFun :: Stmt -> Function
createMainFun stmt =
    let
        evalRayArgs = ["S", "E", "T"]
        evalRgbaArgs = ["X", "D", "T"]
        argTypes = [Float3, Float3, Float]
    in if "X" `elem` varsStmt stmt
        then Function "evalRgba" Float4 (zip evalRgbaArgs argTypes) stmt
        else Function "evalRay" Float3 (zip evalRayArgs argTypes) stmt


convModule :: P.Module a -> Program
-- convModule (P.Module stmts) = removeGroups $ convSuite stmts
convModule (P.Module stmts) =
    let
        Group body = removeGroups $ convSuite stmts
        
        -- extract user function definitions
        f (FunDef _) = True
        f _ = False
        functions = filter f body
        
        -- put the rest into the main function
        body' = filter (Prelude.not . f) body
        mainFun = createMainFun (Group body')
    in [f | FunDef f <- functions] ++ [mainFun]


convSuite :: [P.Statement a] -> Stmt
convSuite s = Group $ map convStatement s


plus x y =  Op "plus" x y
minus x y = Op "minus" x y
mult x y =  Op "mult" x y
div x y =   Op "div" x y
not x =     Func "not" [x]
lt x y =    Func "lt" [x, y]

typeFromName "int"    = Int
typeFromName "float"  = Float
typeFromName "float2" = Float2
typeFromName "float3" = Float3
typeFromName "float4" = Float4


convStatement :: P.Statement a -> Stmt
convStatement stmt = case stmt of
    P.Assign [P.Var (P.Ident x _) _] e _ -> Assign x (convExpr e)
    P.Assign _ _ _ -> error "unsupported left-hand side"
    P.AugmentedAssign (P.Var (P.Ident x _) _) op e _ ->
        case op of
            P.PlusAssign _  -> Assign x (V x `plus` convExpr e)
            P.MinusAssign _ -> Assign x (V x `minus` convExpr e)
            P.MultAssign _  -> Assign x (V x `mult` convExpr e)
            P.DivAssign _   -> Assign x (V x `div` convExpr e)
    P.AugmentedAssign _ _ _ _ -> error "unsupported left-hand side"
    P.Conditional cases elseSuite _ ->
        let
            f [(cond, suite)] elseSuite =
                Branch (convExpr cond) (convSuite suite) (convSuite elseSuite)
            f ((cond, suite) : rest) elseSuite =
                Branch (convExpr cond) (convSuite suite) (f rest elseSuite)
        in f cases elseSuite
    P.Return Nothing _ -> error "need to return a value"
    P.Return (Just e) _ -> Return (convExpr e) NoType
    P.While cond body [] _ -> While (convExpr cond) (convSuite body)
    P.Fun (P.Ident name _) args Nothing body _ ->
        let
            args' = [(x, NoType) | P.Param (P.Ident x _) _ _ _ <- args]
            body' = convSuite body
        in FunDef (Function name NoType args' body')
    P.For [P.Var (P.Ident x _) _] gen body _ _ ->
        case convExpr gen of
            Func "linspace" [eFrom, eTo, eNum] ->
                let
                    cond = V x `lt` eTo
                    body' = convSuite body
                    delta = (eTo `minus` eFrom) `div` eNum
                    incr = Assign x (V x `plus` delta)
                -- in Group [Assign x eFrom, While cond [Group body', incr]]
                 in Group [Assign x eFrom, For cond incr body']
            Func "range" [eNum] ->
                let
                    cond = V x `lt` eNum
                    body' = convSuite body
                    delta = Ci 1
                    incr = Assign x (V x `plus` delta)
                -- in Group [Assign x (Ci 0), While cond [Group body', incr]]
                in Group [Assign x (Ci 0), For cond incr body']
            x -> error $ "unsupported generator in for loop: " ++ show x
    P.Break _ -> Break
    P.Continue _ -> Continue
    P.Pass _ -> Group []
    _ -> error $ "unknown statement: " ++ prettyText stmt
                

convExpr :: P.Expr a -> Expr
convExpr e = case e of
    P.Var (P.Ident "True" _) _ -> T
    P.Var (P.Ident "False" _) _ -> F
    P.Var (P.Ident x _) _ -> V x
    P.Int x _ _ -> Ci x
    P.Float x _ _ -> C x
    P.Bool x _ -> if x then T else F
    P.Strings xs _ -> S (concat xs)
    P.Call f args _ ->
        let 
            trArg (P.ArgExpr e _) = e
            trArg _ = error "only ordinary arguments supported"
            args' = map (convExpr . trArg) args
        in case f of
            P.Var (P.Ident x _) _ ->
                if x `elem` ["min", "max"]
                    then case args' of
                        [e1, e2] -> Op x e1 e2
                        _ -> error $ "wrong number of arguments: " ++ prettyText e
                    else case x of
                        "int" -> Cast Int (head args')
                        "int2" -> case length args' of
                            1 -> Cast Int2 (head args')
                            2 -> Func "make_int2" args'
                            _ -> error $ "wrong number of arguments: " ++ prettyText e
                        "int3" -> case length args' of
                            1 -> Cast Int3 (head args')
                            3 -> Func "make_int3" args'
                            _ -> error $ "wrong number of arguments: " ++ prettyText e
                        "int4" -> case length args' of
                            1 -> Cast Int4 (head args')
                            4 -> Func "make_int4" args'
                            _ -> error $ "wrong number of arguments: " ++ prettyText e
                        "float" -> Cast Float (head args')
                        "float2" -> case length args' of
                            1 -> Cast Float2 (head args')
                            2 -> Func "make_float2" args'
                            _ -> error $ "wrong number of arguments: " ++ prettyText e
                        "float3" -> case length args' of
                            1 -> Cast Float3 (head args')
                            3 -> Func "make_float3" args'
                            _ -> error $ "wrong number of arguments: " ++ prettyText e
                        "float4" -> case length args' of
                            1 -> Cast Float4 (head args')
                            4 -> Func "make_float4" args'
                            _ -> error $ "wrong number of arguments: " ++ prettyText e
                        _ -> Func x args'
            _ -> error "function name expected"
    P.BinaryOp op e1 e2 _ ->
        let
            e1' = convExpr e1
            e2' = convExpr e2
        in case op of
            P.Dot _ -> Dot e1' x where V x = e2'
            P.And _ -> Func "and" [e1', e2']
            P.Or _ -> Func "or" [e1', e2']
            P.Equality _ -> Func "eq" [e1', e2']
            P.NotEquals _ -> not (Func "eq" [e1', e2'])
            P.GreaterThan _ -> Func "gt" [e1', e2']
            P.GreaterThanEquals _ -> not (Func "lt" [e1', e2'])
            P.LessThan _ -> Func "lt" [e1', e2']
            P.LessThanEquals _ -> not (Func "gt" [e1', e2'])
            P.Plus _ -> Op "plus" e1' e2'
            P.Minus _ -> Op "minus" e1' e2'
            P.Multiply _ -> Op "mult" e1' e2'
            P.Divide _ -> Op "div" e1' e2'
            P.Exponent _ -> Op "pow" e1' e2'
    P.UnaryOp (P.Minus _) e _ ->
        case e of
            P.Int x _ _ -> Ci (-x)
            P.Float x _ _ -> C (-x)
            e' -> (Ci 0) `minus` convExpr e'
    P.UnaryOp (P.Not _) e _ -> Func "not" [convExpr e]
    P.Paren e _ -> convExpr e
    _ -> error $ "Unknown expression: " ++ prettyText e




